"""
技术分析图表模块 - 提供技术分析图表生成功能
"""

import pandas as pd
import numpy as np
import plotly.graph_objects as go
from typing import List, Dict, Any, Optional, Tuple

def create_technical_chart(
    df: pd.DataFrame, 
    title: str = "技术分析", 
    height: int = 600,
    indicators: List[str] = ['MA', 'MACD', 'KDJ', 'RSI']
) -> go.Figure:
    """
    创建技术分析图表
    
    参数:
        df: 包含OHLC数据的DataFrame，必须包含 '日期', '开盘', '最高', '最低', '收盘', '成交量' 列
        title: 图表标题
        height: 图表高度
        indicators: 要显示的技术指标列表
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = ['日期', '开盘', '最高', '最低', '收盘']
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 计算技术指标
    df_with_indicators = calculate_indicators(df, indicators)
    
    # 创建子图布局
    fig = make_subplots(df_with_indicators, indicators, title, height)
    
    return fig

def calculate_indicators(df: pd.DataFrame, indicators: List[str]) -> pd.DataFrame:
    """
    计算技术指标
    
    参数:
        df: 原始数据DataFrame
        indicators: 要计算的技术指标列表
    
    返回:
        添加了技术指标的DataFrame
    """
    # 创建DataFrame的副本，避免修改原始数据
    result_df = df.copy()
    
    # 计算移动平均线
    if 'MA' in indicators:
        result_df['MA5'] = result_df['收盘'].rolling(window=5).mean()
        result_df['MA10'] = result_df['收盘'].rolling(window=10).mean()
        result_df['MA20'] = result_df['收盘'].rolling(window=20).mean()
        result_df['MA60'] = result_df['收盘'].rolling(window=60).mean()
    
    # 计算MACD
    if 'MACD' in indicators:
        # 计算EMA
        result_df['EMA12'] = result_df['收盘'].ewm(span=12, adjust=False).mean()
        result_df['EMA26'] = result_df['收盘'].ewm(span=26, adjust=False).mean()
        
        # 计算DIF和DEA
        result_df['DIF'] = result_df['EMA12'] - result_df['EMA26']
        result_df['DEA'] = result_df['DIF'].ewm(span=9, adjust=False).mean()
        
        # 计算MACD柱状
        result_df['MACD'] = 2 * (result_df['DIF'] - result_df['DEA'])
    
    # 计算KDJ
    if 'KDJ' in indicators:
        # 计算RSV
        low_9 = result_df['最低'].rolling(window=9).min()
        high_9 = result_df['最高'].rolling(window=9).max()
        rsv = 100 * ((result_df['收盘'] - low_9) / (high_9 - low_9))
        
        # 计算K、D、J值
        result_df['K'] = rsv.ewm(alpha=1/3, adjust=False).mean()
        result_df['D'] = result_df['K'].ewm(alpha=1/3, adjust=False).mean()
        result_df['J'] = 3 * result_df['K'] - 2 * result_df['D']
    
    # 计算RSI
    if 'RSI' in indicators:
        # 计算价格变化
        delta = result_df['收盘'].diff()
        
        # 分离上涨和下跌
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        
        # 计算平均上涨和下跌
        avg_gain_14 = gain.rolling(window=14).mean()
        avg_loss_14 = loss.rolling(window=14).mean()
        
        # 计算相对强度
        rs_14 = avg_gain_14 / avg_loss_14
        
        # 计算RSI
        result_df['RSI14'] = 100 - (100 / (1 + rs_14))
    
    # 计算布林带
    if 'BOLL' in indicators:
        # 计算中轨（20日移动平均线）
        result_df['BOLL_MID'] = result_df['收盘'].rolling(window=20).mean()
        
        # 计算标准差
        result_df['BOLL_STD'] = result_df['收盘'].rolling(window=20).std()
        
        # 计算上轨和下轨
        result_df['BOLL_UP'] = result_df['BOLL_MID'] + 2 * result_df['BOLL_STD']
        result_df['BOLL_DOWN'] = result_df['BOLL_MID'] - 2 * result_df['BOLL_STD']
    
    return result_df

def make_subplots(df: pd.DataFrame, indicators: List[str], title: str, height: int) -> go.Figure:
    """
    创建子图布局
    
    参数:
        df: 包含技术指标的DataFrame
        indicators: 要显示的技术指标列表
        title: 图表标题
        height: 图表高度
    
    返回:
        plotly图表对象
    """
    # 计算需要的子图数量
    num_subplots = 1  # 主K线图
    
    if 'MACD' in indicators:
        num_subplots += 1
    
    if 'KDJ' in indicators:
        num_subplots += 1
    
    if 'RSI' in indicators:
        num_subplots += 1
    
    if '成交量' in df.columns:
        num_subplots += 1
    
    # 创建子图
    fig = go.Figure()
    
    # 添加K线图
    fig.add_trace(go.Candlestick(
        x=df['日期'],
        open=df['开盘'],
        high=df['最高'],
        low=df['最低'],
        close=df['收盘'],
        name='K线',
        increasing_line_color='red',
        decreasing_line_color='green'
    ))
    
    # 添加移动平均线
    if 'MA' in indicators:
        ma_colors = ['blue', 'orange', 'purple', 'cyan']
        ma_periods = [5, 10, 20, 60]
        
        for i, period in enumerate(ma_periods):
            ma_col = f'MA{period}'
            if ma_col in df.columns:
                fig.add_trace(go.Scatter(
                    x=df['日期'],
                    y=df[ma_col],
                    mode='lines',
                    name=ma_col,
                    line=dict(color=ma_colors[i])
                ))
    
    # 添加布林带
    if 'BOLL' in indicators and 'BOLL_MID' in df.columns:
        fig.add_trace(go.Scatter(
            x=df['日期'],
            y=df['BOLL_UP'],
            mode='lines',
            name='BOLL上轨',
            line=dict(color='rgba(250, 128, 114, 0.7)', dash='dash')
        ))
        
        fig.add_trace(go.Scatter(
            x=df['日期'],
            y=df['BOLL_MID'],
            mode='lines',
            name='BOLL中轨',
            line=dict(color='rgba(70, 130, 180, 0.7)')
        ))
        
        fig.add_trace(go.Scatter(
            x=df['日期'],
            y=df['BOLL_DOWN'],
            mode='lines',
            name='BOLL下轨',
            line=dict(color='rgba(152, 251, 152, 0.7)', dash='dash')
        ))
    
    # 设置主图布局
    fig.update_layout(
        title=title,
        xaxis_title="日期",
        yaxis_title="价格",
        height=height,
        xaxis_rangeslider_visible=False,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    # 这里简化处理，实际上应该使用make_subplots创建多个子图
    # 但由于复杂度较高，这里只展示主图
    
    return fig

def create_support_resistance_chart(
    df: pd.DataFrame, 
    support_levels: List[float],
    resistance_levels: List[float],
    title: str = "支撑与阻力位", 
    height: int = 600
) -> go.Figure:
    """
    创建支撑与阻力位图表
    
    参数:
        df: 包含OHLC数据的DataFrame，必须包含 '日期', '开盘', '最高', '最低', '收盘' 列
        support_levels: 支撑位列表
        resistance_levels: 阻力位列表
        title: 图表标题
        height: 图表高度
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = ['日期', '开盘', '最高', '最低', '收盘']
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 创建图表
    fig = go.Figure()
    
    # 添加K线图
    fig.add_trace(go.Candlestick(
        x=df['日期'],
        open=df['开盘'],
        high=df['最高'],
        low=df['最低'],
        close=df['收盘'],
        name='K线',
        increasing_line_color='red',
        decreasing_line_color='green'
    ))
    
    # 添加支撑位
    for level in support_levels:
        fig.add_shape(
            type="line",
            x0=df['日期'].iloc[0],
            y0=level,
            x1=df['日期'].iloc[-1],
            y1=level,
            line=dict(
                color="green",
                width=2,
                dash="dash",
            ),
            name=f"支撑位 {level}"
        )
        
        # 添加文本标签
        fig.add_annotation(
            x=df['日期'].iloc[-1],
            y=level,
            text=f"支撑位: {level}",
            showarrow=False,
            yshift=10,
            bgcolor="rgba(0,255,0,0.3)"
        )
    
    # 添加阻力位
    for level in resistance_levels:
        fig.add_shape(
            type="line",
            x0=df['日期'].iloc[0],
            y0=level,
            x1=df['日期'].iloc[-1],
            y1=level,
            line=dict(
                color="red",
                width=2,
                dash="dash",
            ),
            name=f"阻力位 {level}"
        )
        
        # 添加文本标签
        fig.add_annotation(
            x=df['日期'].iloc[-1],
            y=level,
            text=f"阻力位: {level}",
            showarrow=False,
            yshift=-20,
            bgcolor="rgba(255,0,0,0.3)"
        )
    
    # 设置图表布局
    fig.update_layout(
        title=title,
        xaxis_title="日期",
        yaxis_title="价格",
        height=height,
        xaxis_rangeslider_visible=False
    )
    
    return fig